{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Network with SPU\n",
    "\n",
    ">  Please read lab `Logistic Regression On SPU` first if you have not。\n",
    "\n",
    "In lab `Logistic Regression On SPU`, we have showed how to use SecretFlow/SPU to convert a plaintext JAX training program to a secrue MPC trainning program.\n",
    "\n",
    "In this lab, the idea is quite similar but this time we will work with a Neural Network model.\n",
    "\n",
    "We are going to use the same dataset and all the settings as lab `Logistic Regression On SPU`.\n",
    "\n",
    "And first, let's work out the plaintext model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*The following codes are demos only. It's **NOT for production** due to system security concerns, please **DO NOT** use it directly in production.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a model with JAX/FLAX\n",
    "\n",
    "### Load the Dataset\n",
    "\n",
    "The below is just copied from lab `Logistic Regression On SPU`. I'm not going to explain again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.datasets import load_breast_cancer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "\n",
    "def load_train_dataset(party_id=None) -> (np.ndarray, np.ndarray):\n",
    "    features, label = load_breast_cancer(return_X_y=True)\n",
    "    scaler = StandardScaler()\n",
    "    features = scaler.fit_transform(features)\n",
    "    X_train, _, y_train, _ = train_test_split(\n",
    "        features, label, test_size=0.8, random_state=42\n",
    "    )\n",
    "\n",
    "    if party_id:\n",
    "        if party_id == 1:\n",
    "            return X_train[:, 15:], _\n",
    "        else:\n",
    "            return X_train[:, :15], y_train\n",
    "    else:\n",
    "        return X_train, y_train\n",
    "\n",
    "\n",
    "def load_test_dataset():\n",
    "    features, label = load_breast_cancer(return_X_y=True)\n",
    "    scaler = StandardScaler()\n",
    "    features = scaler.fit_transform(features)\n",
    "    _, X_test, _, y_test = train_test_split(\n",
    "        features, label, test_size=0.8, random_state=42\n",
    "    )\n",
    "    return X_test, y_test\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Model\n",
    "\n",
    "\n",
    "We are going to use a 4-layer [MLP](https://en.wikipedia.org/wiki/Multilayer_perceptron) model with a [ReLU](https://en.wikipedia.org/wiki/Rectifier_(neural_networks)) activation function here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.example_libraries import stax\n",
    "from jax.example_libraries.stax import (\n",
    "    Dense,\n",
    "    Relu,\n",
    ")\n",
    "\n",
    "\n",
    "def MLP():\n",
    "    nn_init, nn_apply = stax.serial(\n",
    "        Dense(30),\n",
    "        Relu,\n",
    "        Dense(15),\n",
    "        Relu,\n",
    "        Dense(8),\n",
    "        Relu,\n",
    "        Dense(1),\n",
    "    )\n",
    "\n",
    "    return nn_init, nn_apply\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we define the trainning method here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax.example_libraries import optimizers, stax\n",
    "\n",
    "\n",
    "KEY = jax.random.PRNGKey(42)\n",
    "INPUT_SHAPE = (-1,30)\n",
    "\n",
    "def init_state(learning_rate):\n",
    "    init_fun, _ = MLP()\n",
    "    _, params_init = init_fun(KEY, INPUT_SHAPE)\n",
    "    opt_init, _, _ = optimizers.sgd(learning_rate)\n",
    "    opt_state = opt_init(params_init)\n",
    "    return opt_state\n",
    "\n",
    "def train(\n",
    "    train_x1,\n",
    "    train_x2,\n",
    "    train_y,\n",
    "    opt_state,\n",
    "    learning_rate,\n",
    "    epochs,\n",
    "    batch_size,\n",
    "):\n",
    "    train_x = jnp.concatenate([train_x1, train_x2], axis=1)\n",
    "\n",
    "    _, predict_fun = MLP()\n",
    "    _, opt_update, get_params = optimizers.sgd(learning_rate)\n",
    "\n",
    "    def update_model(state, imgs, labels, i):\n",
    "        def mse(y, pred):\n",
    "            return jnp.mean(jnp.multiply(y - pred, y - pred) / 2.0)\n",
    "\n",
    "        def loss_func(params):\n",
    "            y = predict_fun(params, imgs)\n",
    "            return mse(y, labels), y\n",
    "\n",
    "        grad_fn = jax.value_and_grad(loss_func, has_aux=True)\n",
    "        (loss, y), grads = grad_fn(get_params(state))\n",
    "        return opt_update(i, grads, state)\n",
    "\n",
    "    for i in range(1, epochs + 1):\n",
    "        imgs_batchs = jnp.array_split(train_x, len(train_x) / batch_size, axis=0)\n",
    "        labels_batchs = jnp.array_split(train_y, len(train_y) / batch_size, axis=0)\n",
    "\n",
    "        for batch_idx, (batch_images, batch_labels) in enumerate(\n",
    "            zip(imgs_batchs, labels_batchs)\n",
    "        ):\n",
    "            opt_state = update_model(opt_state, batch_images, batch_labels, i)\n",
    "    return get_params(opt_state)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validate the Model\n",
    "\n",
    "We use AUC as the validation metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "\n",
    "def validate_model(params, X_test, y_test):\n",
    "    _, predict_fun = MLP()\n",
    "    y_pred = predict_fun(params, X_test)\n",
    "    return roc_auc_score(y_test, y_pred)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BUILD Together\n",
    "\n",
    "Let's put everything together and train a plaintext NN model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(DeviceArray([[ 0.05228749,  0.2254616 ,  0.14233142,  0.2218795 ,\n",
      "               0.14313836, -0.11788452, -0.25514612,  0.05024785,\n",
      "               0.19256376, -0.16702715,  0.1360347 ,  0.09218568,\n",
      "              -0.14514978, -0.00203137,  0.1026915 ,  0.16864221,\n",
      "              -0.19459064, -0.1680608 , -0.0291565 ,  0.24780752,\n",
      "              -0.10630313,  0.19863537, -0.07401107, -0.24004339,\n",
      "              -0.11391521,  0.3779385 ,  0.1651407 ,  0.0792419 ,\n",
      "              -0.21970308, -0.07381143],\n",
      "             [-0.23652235,  0.3905772 , -0.0325642 ,  0.16153163,\n",
      "               0.03559006, -0.34141707,  0.36782634,  0.35197827,\n",
      "              -0.03010983,  0.17985871,  0.1022082 ,  0.15198684,\n",
      "               0.20280722, -0.15535937, -0.3000578 ,  0.28659174,\n",
      "              -0.2538568 , -0.10807706,  0.02643047, -0.01051225,\n",
      "              -0.23814394, -0.11327465, -0.2019997 ,  0.03128352,\n",
      "              -0.046227  ,  0.13949452, -0.05219001,  0.10902884,\n",
      "               0.134705  , -0.11092291],\n",
      "             [-0.18260652,  0.16455476,  0.06864081,  0.09414306,\n",
      "              -0.18830639, -0.08217652, -0.13038777, -0.4105987 ,\n",
      "               0.23539643,  0.1685532 ,  0.04342773,  0.0273842 ,\n",
      "               0.1272156 ,  0.01272896, -0.19060796, -0.04190553,\n",
      "              -0.26819235,  0.07810482,  0.03015172,  0.15575643,\n",
      "              -0.07964827, -0.30778813, -0.03847689, -0.35754323,\n",
      "              -0.1945856 ,  0.04414607, -0.22458269, -0.19626944,\n",
      "               0.12733065, -0.14534126],\n",
      "             [-0.2133216 ,  0.19377995,  0.04313114,  0.16396867,\n",
      "               0.09944846,  0.21389818,  0.22278464,  0.0648506 ,\n",
      "              -0.03629925,  0.12235434, -0.06149251,  0.07482152,\n",
      "               0.08981636, -0.16575499, -0.08681874,  0.32138425,\n",
      "              -0.14770229, -0.01284664,  0.08067241, -0.25103945,\n",
      "               0.14798449, -0.40178353, -0.1300967 , -0.3074095 ,\n",
      "               0.15551645,  0.05096855, -0.38844764, -0.02998807,\n",
      "              -0.17790447, -0.33897516],\n",
      "             [ 0.2845987 , -0.13556221,  0.11223896, -0.11234169,\n",
      "               0.24697033, -0.12554286, -0.00190957, -0.05253758,\n",
      "              -0.0847037 , -0.05462934, -0.25783134,  0.24276951,\n",
      "              -0.11871891,  0.2850348 , -0.1303706 , -0.06938016,\n",
      "               0.15665027,  0.06147627, -0.05267749,  0.00233302,\n",
      "               0.04734558,  0.1760517 , -0.03966657,  0.11916655,\n",
      "              -0.25819635, -0.32799518, -0.21394247,  0.2564436 ,\n",
      "              -0.34983522, -0.06176936],\n",
      "             [-0.17650877,  0.19584092,  0.25888008,  0.1280048 ,\n",
      "               0.07615547,  0.04855842, -0.14355224,  0.01383955,\n",
      "               0.12403549, -0.07664248, -0.00134847, -0.06848545,\n",
      "              -0.3606049 ,  0.06312151, -0.13296399,  0.30586514,\n",
      "               0.19914857,  0.03221894,  0.24011207, -0.29851496,\n",
      "               0.02360805,  0.00451724, -0.06560124, -0.11889942,\n",
      "               0.13419034,  0.3323109 , -0.30779368,  0.12421666,\n",
      "              -0.05408807,  0.0449027 ],\n",
      "             [-0.24534445, -0.17287417, -0.15903771,  0.11636478,\n",
      "              -0.24853751, -0.2411471 , -0.03998812,  0.36924326,\n",
      "               0.06296232, -0.05439488,  0.2099218 , -0.3997368 ,\n",
      "              -0.14747825,  0.06252322, -0.3339467 , -0.15357825,\n",
      "              -0.22988296, -0.13085324,  0.05122835, -0.03729573,\n",
      "               0.12808837,  0.109161  , -0.09478054,  0.09113201,\n",
      "              -0.12293261, -0.03382579,  0.2969238 ,  0.17922421,\n",
      "              -0.03303223,  0.03031438],\n",
      "             [-0.13064878, -0.12434947, -0.03403151,  0.34820265,\n",
      "               0.05029985,  0.09225762,  0.3157716 , -0.16092533,\n",
      "              -0.08145321, -0.05226815,  0.16237862,  0.32751697,\n",
      "              -0.01519725, -0.16877326, -0.32174802,  0.34621912,\n",
      "               0.25516868,  0.22436927,  0.21228592,  0.21658626,\n",
      "               0.03310975,  0.22128862,  0.37305608,  0.16911261,\n",
      "               0.02211326, -0.16129407,  0.14026178,  0.04060666,\n",
      "              -0.16443706,  0.06357209],\n",
      "             [ 0.07629905,  0.18016288,  0.20691985, -0.02306781,\n",
      "              -0.05998988,  0.3135441 , -0.07086642,  0.05441632,\n",
      "              -0.04780589,  0.23343606,  0.3195129 , -0.24244739,\n",
      "              -0.15932274,  0.08182976, -0.08398597, -0.17806613,\n",
      "               0.27671754, -0.02641992,  0.24031332, -0.14868899,\n",
      "               0.1358432 ,  0.15450896,  0.33141395, -0.07921614,\n",
      "              -0.01936086,  0.15076634,  0.02756989,  0.29221758,\n",
      "               0.12562875, -0.13665336],\n",
      "             [-0.15456939,  0.2587808 , -0.3999844 ,  0.13280822,\n",
      "              -0.05051823, -0.1336271 , -0.332645  ,  0.05919369,\n",
      "              -0.10773447, -0.08454138, -0.1097419 ,  0.07605839,\n",
      "              -0.12866594,  0.22350465,  0.12161157, -0.02826402,\n",
      "              -0.28843305,  0.06870433, -0.0702437 , -0.15990552,\n",
      "              -0.1616978 , -0.14658587,  0.0637271 ,  0.01703468,\n",
      "               0.14891641,  0.04488591, -0.00506696, -0.31588253,\n",
      "              -0.04548473, -0.00550163],\n",
      "             [ 0.24171245,  0.07552842,  0.29506242,  0.1468854 ,\n",
      "               0.09751622,  0.22560577, -0.06077249, -0.1744719 ,\n",
      "               0.03879349,  0.2177638 , -0.09062548,  0.0885776 ,\n",
      "               0.13792096,  0.41941765,  0.07514104, -0.21747787,\n",
      "               0.16517821, -0.00159747, -0.09909979,  0.14524716,\n",
      "              -0.11354675,  0.05275302,  0.12992746,  0.38263318,\n",
      "              -0.31539124,  0.18909855, -0.13415688,  0.13672829,\n",
      "               0.02819634, -0.1723328 ],\n",
      "             [-0.02140586,  0.04052684, -0.08711036,  0.06792393,\n",
      "               0.09729451,  0.06603958, -0.20308013, -0.11747426,\n",
      "               0.16307876,  0.2781937 , -0.30456808, -0.02802942,\n",
      "              -0.16799237,  0.3521512 , -0.05961246,  0.14780895,\n",
      "              -0.2782957 , -0.08539459, -0.00634822, -0.11484689,\n",
      "               0.00630033,  0.09019561, -0.11777002,  0.11413086,\n",
      "              -0.17090258,  0.15509939, -0.03246784, -0.34619752,\n",
      "               0.0654332 ,  0.18084455],\n",
      "             [ 0.07866376,  0.02720631, -0.1249831 ,  0.22738136,\n",
      "               0.02853629,  0.0281287 , -0.23032783,  0.22130139,\n",
      "              -0.03210012, -0.17050035, -0.22716315, -0.2739657 ,\n",
      "              -0.12097008,  0.13619848, -0.10620885, -0.06121876,\n",
      "              -0.05841545, -0.25739062,  0.3592187 , -0.03755495,\n",
      "               0.04315823,  0.07351541, -0.36266065, -0.04485212,\n",
      "               0.12890156,  0.30077755, -0.18278316, -0.14822567,\n",
      "              -0.36666572, -0.3292698 ],\n",
      "             [ 0.00287816,  0.01815061, -0.22103727, -0.00184069,\n",
      "              -0.12475448,  0.16648997,  0.10488272,  0.24595365,\n",
      "               0.33036956, -0.31138647,  0.38887528, -0.2326023 ,\n",
      "               0.05298411, -0.04623241,  0.06839247,  0.3464936 ,\n",
      "              -0.20316567,  0.36484316, -0.28965023,  0.01836067,\n",
      "               0.32735577,  0.245742  ,  0.04199012,  0.03530597,\n",
      "              -0.02775852, -0.28184822,  0.20460907, -0.30127418,\n",
      "               0.03575421,  0.40105996],\n",
      "             [ 0.28987   ,  0.19460946, -0.07549638, -0.18909688,\n",
      "               0.0952258 , -0.16651033,  0.02105147, -0.09171448,\n",
      "              -0.2534539 , -0.01885388, -0.01193424,  0.04423714,\n",
      "               0.10957579,  0.38906747, -0.04517176, -0.04878805,\n",
      "              -0.31367528, -0.01723027,  0.01444988, -0.00301655,\n",
      "              -0.27081043, -0.22861753, -0.09883647, -0.06335558,\n",
      "              -0.15288207,  0.02695066, -0.03355371,  0.01292902,\n",
      "              -0.00920994,  0.04239657],\n",
      "             [ 0.03434585,  0.15078123,  0.0827676 , -0.17750008,\n",
      "               0.02391125, -0.33295327, -0.00113884,  0.19944322,\n",
      "               0.33309713, -0.10293391, -0.38618195,  0.02625184,\n",
      "              -0.24926597,  0.05332365, -0.05411308, -0.25822756,\n",
      "              -0.13695747,  0.16165987,  0.1154609 , -0.30235574,\n",
      "              -0.07970641, -0.19536158,  0.31444794,  0.24430597,\n",
      "              -0.07022712,  0.40057224,  0.10494798,  0.01605886,\n",
      "              -0.16881308,  0.03208243],\n",
      "             [-0.08846725,  0.15857734, -0.24452923,  0.02802754,\n",
      "              -0.28953943,  0.31910077, -0.01349073,  0.150713  ,\n",
      "              -0.12966292, -0.18281364, -0.19632503,  0.09208366,\n",
      "              -0.17555399,  0.16875768,  0.0425629 , -0.31164765,\n",
      "               0.25497162,  0.12969662, -0.06240876,  0.20764005,\n",
      "              -0.01605446, -0.08775843,  0.24665995, -0.01994762,\n",
      "               0.01353364,  0.08691632,  0.17876448,  0.263216  ,\n",
      "              -0.12575743, -0.28093317],\n",
      "             [-0.27248427,  0.39057368,  0.15252772,  0.00340817,\n",
      "              -0.17826745, -0.22753881, -0.2125268 , -0.01392176,\n",
      "               0.24692835,  0.14416543, -0.276925  ,  0.29618964,\n",
      "               0.04703401, -0.0887985 , -0.0971118 , -0.02631477,\n",
      "              -0.02944293,  0.18631102,  0.20172389,  0.12907964,\n",
      "              -0.14135914, -0.2149143 ,  0.09168746, -0.04743837,\n",
      "              -0.20446979, -0.24302764,  0.01107965,  0.0774961 ,\n",
      "               0.01269388,  0.0724828 ],\n",
      "             [-0.08600874,  0.07436551, -0.08783555, -0.0167344 ,\n",
      "              -0.34101936,  0.03360803, -0.03720398,  0.3244064 ,\n",
      "              -0.0566737 , -0.22823003, -0.12521851,  0.21265425,\n",
      "              -0.01466573,  0.12385557, -0.04733623, -0.0135651 ,\n",
      "               0.20629783, -0.0052271 , -0.12534957,  0.06386559,\n",
      "               0.20566939,  0.19607843,  0.41228947,  0.24781   ,\n",
      "              -0.2929769 , -0.23787865,  0.01076282,  0.24403125,\n",
      "               0.06110474,  0.02741824],\n",
      "             [ 0.15478764, -0.06889063,  0.02127939, -0.14462551,\n",
      "               0.06981816,  0.03925344, -0.1508263 ,  0.31153485,\n",
      "               0.16261367,  0.0781375 ,  0.08291113,  0.1948986 ,\n",
      "              -0.03388112,  0.00376782, -0.00776635, -0.37110853,\n",
      "               0.18846121,  0.13249442, -0.02163309, -0.09489319,\n",
      "              -0.13345313, -0.02784892,  0.26336747, -0.1795453 ,\n",
      "               0.21551102, -0.08535276,  0.06097673, -0.04588852,\n",
      "               0.13760102, -0.1133495 ],\n",
      "             [-0.11020133, -0.23387879, -0.12427491,  0.06956701,\n",
      "               0.01823351, -0.17487676,  0.02797757,  0.11959445,\n",
      "              -0.27016315,  0.19566973,  0.19247875,  0.25272843,\n",
      "              -0.13282701,  0.06479955, -0.09028956,  0.07406615,\n",
      "               0.00682314, -0.03973693,  0.17775624, -0.23105612,\n",
      "               0.35531384,  0.2822218 , -0.17894082,  0.06605034,\n",
      "               0.12765801,  0.11652758,  0.13850042,  0.12908229,\n",
      "               0.33359197,  0.1672675 ],\n",
      "             [ 0.19818322,  0.08144253, -0.09238429, -0.11318122,\n",
      "              -0.05305329,  0.08116519, -0.39225355, -0.10105433,\n",
      "               0.21901254, -0.17041683,  0.27980962, -0.11357052,\n",
      "              -0.1519309 , -0.01640943,  0.18932036, -0.28740293,\n",
      "               0.20391062,  0.18099917,  0.14320657,  0.18867284,\n",
      "              -0.17415746,  0.31392294, -0.06469122, -0.24698544,\n",
      "              -0.18791103,  0.07078048,  0.07739505,  0.05967075,\n",
      "              -0.18496847,  0.05231325],\n",
      "             [ 0.37151197, -0.02902467, -0.18655425, -0.01155737,\n",
      "              -0.15688387, -0.11883423,  0.04187197,  0.43013635,\n",
      "               0.05934869, -0.11776802,  0.1501778 ,  0.16188078,\n",
      "               0.0526493 ,  0.23144269, -0.19149637,  0.0946878 ,\n",
      "              -0.01882421,  0.23430654,  0.22797754,  0.18076888,\n",
      "               0.00203844, -0.10573798, -0.11869374,  0.14509451,\n",
      "              -0.04800387,  0.0588996 , -0.33298695,  0.05973347,\n",
      "               0.22180122,  0.15065655],\n",
      "             [-0.13513307, -0.12206829, -0.13369805, -0.10280453,\n",
      "              -0.17875108, -0.39313182,  0.38611713,  0.00054136,\n",
      "              -0.06609095, -0.2119039 , -0.03643068, -0.10959922,\n",
      "               0.21848267, -0.0718011 ,  0.05998127, -0.01975773,\n",
      "               0.03837547,  0.25704667,  0.0869022 ,  0.12466951,\n",
      "               0.03259438,  0.00690077, -0.10029537, -0.32062462,\n",
      "              -0.06477929, -0.03771692, -0.1000107 ,  0.3223988 ,\n",
      "               0.18602814, -0.26882583],\n",
      "             [-0.07199647,  0.09197637, -0.3337017 , -0.32504794,\n",
      "              -0.21852113,  0.16968639,  0.22908044, -0.1121565 ,\n",
      "               0.00849434, -0.23621182,  0.24395975, -0.19499105,\n",
      "               0.09914723,  0.08460287, -0.17264386,  0.04430242,\n",
      "              -0.12306613,  0.03333357, -0.3071449 ,  0.14468828,\n",
      "              -0.17336847, -0.12829198, -0.20469593,  0.08525001,\n",
      "               0.32388493,  0.14959903, -0.01399714,  0.1650289 ,\n",
      "               0.25087044,  0.05552226],\n",
      "             [ 0.07694902, -0.15461436, -0.13562253, -0.03767516,\n",
      "               0.2882856 , -0.21340683,  0.31948557,  0.21057953,\n",
      "               0.15224065, -0.21047139,  0.11142458,  0.11242549,\n",
      "               0.26879036,  0.17542458,  0.00410767, -0.0238772 ,\n",
      "              -0.1730145 , -0.06978362,  0.05905372, -0.30531165,\n",
      "               0.02002045, -0.09711721, -0.10325058, -0.02199581,\n",
      "               0.26327002,  0.03562784, -0.05885479, -0.19802903,\n",
      "               0.25269374,  0.06796943],\n",
      "             [ 0.36213663, -0.03719085,  0.2661024 ,  0.06537861,\n",
      "               0.15390968, -0.02520406,  0.02858835, -0.13868803,\n",
      "               0.1566722 , -0.24320574, -0.18535204,  0.07182785,\n",
      "               0.16526997, -0.10515139, -0.1586223 , -0.21504097,\n",
      "               0.08026339, -0.2787892 , -0.02720969, -0.05869072,\n",
      "              -0.0102096 ,  0.2692256 ,  0.37070578,  0.15374315,\n",
      "              -0.20913827, -0.06358949,  0.16481334, -0.15972714,\n",
      "              -0.09043851, -0.31333506],\n",
      "             [-0.26683992, -0.2932892 , -0.37518927,  0.05802228,\n",
      "              -0.05382813,  0.0516159 , -0.19448185,  0.27751648,\n",
      "               0.10662734,  0.02933954,  0.13876936,  0.15566003,\n",
      "               0.09799541, -0.26987612,  0.21156615, -0.15326263,\n",
      "               0.2841693 ,  0.2547364 ,  0.03386704, -0.19351305,\n",
      "              -0.17508656, -0.08875425,  0.16013508, -0.09308628,\n",
      "               0.20142552, -0.19189563, -0.10587319,  0.09671734,\n",
      "               0.05608173,  0.03757936],\n",
      "             [-0.04800365, -0.07009161,  0.31190217, -0.24311927,\n",
      "               0.37147722, -0.31526932,  0.17272617, -0.21305798,\n",
      "              -0.31096208,  0.19638465,  0.23543732, -0.16582349,\n",
      "               0.00981798, -0.12924755, -0.07474051,  0.11896998,\n",
      "              -0.20260325,  0.13911168,  0.1547592 , -0.00315949,\n",
      "               0.08139693,  0.01813293,  0.281589  ,  0.13167155,\n",
      "              -0.22986083,  0.3450412 ,  0.26744112,  0.2590702 ,\n",
      "               0.15566453, -0.0023672 ],\n",
      "             [ 0.00611473,  0.0605232 ,  0.0747226 , -0.0231379 ,\n",
      "              -0.12747358, -0.12102512, -0.1129632 , -0.31989118,\n",
      "               0.13601954, -0.22598006,  0.08662641,  0.38254005,\n",
      "               0.23247495,  0.17463045,  0.04016639, -0.10817838,\n",
      "              -0.15106298, -0.20628011, -0.26059455,  0.10026004,\n",
      "               0.2203435 ,  0.246053  , -0.28631967,  0.16177788,\n",
      "               0.08512966,  0.1777773 , -0.18250732,  0.05392032,\n",
      "              -0.00266484,  0.01872227]], dtype=float32), DeviceArray([ 0.01719807,  0.01454616,  0.01905813,  0.01762039,\n",
      "             -0.0013846 , -0.02471952, -0.00305147,  0.02690976,\n",
      "             -0.02910551, -0.0295884 , -0.01861479,  0.00952155,\n",
      "             -0.03130547,  0.02270429,  0.01055896, -0.04029678,\n",
      "              0.02134624,  0.00379404,  0.00251089, -0.00785888,\n",
      "             -0.02293164, -0.02042793,  0.03010432, -0.0206696 ,\n",
      "             -0.01038983, -0.0005016 , -0.03257501,  0.01629826,\n",
      "              0.00870296, -0.00119909], dtype=float32)), (), (DeviceArray([[ 2.38768652e-01,  1.11299731e-01, -6.66937381e-02,\n",
      "              -1.53782055e-01, -1.61151849e-02, -2.14111865e-01,\n",
      "               4.11888123e-01, -1.79352835e-01, -9.40110981e-02,\n",
      "              -2.57594913e-01, -2.79346630e-02, -2.03594118e-01,\n",
      "              -1.39336213e-01,  2.16380909e-01,  1.99046016e-01],\n",
      "             [-2.03276083e-01,  2.72610728e-02,  2.14246139e-02,\n",
      "               1.91209987e-01, -9.25622508e-02, -1.01939410e-01,\n",
      "              -1.00079384e-02, -1.39731899e-01,  2.97634065e-01,\n",
      "              -5.27764000e-02, -1.32999450e-01, -1.73676386e-01,\n",
      "              -2.68746428e-02,  3.85802418e-01,  3.27703953e-01],\n",
      "             [-2.27078702e-02, -2.26532757e-01, -3.49178016e-02,\n",
      "              -2.46451825e-01, -9.35519636e-02, -6.53570071e-02,\n",
      "               3.49297049e-03, -3.24385077e-01,  3.30161937e-02,\n",
      "               6.76832944e-02, -2.00153496e-02, -1.61580145e-01,\n",
      "               3.19239289e-01, -1.99406728e-01,  5.22819161e-02],\n",
      "             [ 1.89090952e-01, -2.82470971e-01, -2.34250307e-01,\n",
      "               2.82150000e-01, -2.85152525e-01, -8.81261528e-02,\n",
      "              -4.46520478e-01,  2.65962452e-01, -1.03820488e-01,\n",
      "               1.94392309e-01,  1.69514671e-01, -8.82153884e-02,\n",
      "              -1.13377139e-01, -8.65748599e-02, -4.72750217e-02],\n",
      "             [-4.34734583e-01,  3.04439978e-04, -2.06813246e-01,\n",
      "              -3.43766436e-02, -9.85913575e-02, -1.30387889e-02,\n",
      "              -1.78572848e-01,  1.38609096e-01,  2.01166961e-02,\n",
      "              -6.78521097e-02, -1.91368133e-01,  1.01623885e-01,\n",
      "              -1.84639305e-01,  4.62555066e-02,  1.94816981e-02],\n",
      "             [ 3.88318226e-02,  2.38659933e-01,  7.73253441e-02,\n",
      "              -4.95840311e-02,  3.21680933e-01,  2.42353261e-01,\n",
      "               1.17629588e-01,  3.30427736e-01,  2.49388218e-02,\n",
      "              -2.39334166e-01,  1.80144161e-01, -2.65455157e-01,\n",
      "              -1.70674339e-01,  1.92125559e-01,  7.25867152e-02],\n",
      "             [-1.40786320e-01,  4.83135134e-03,  2.86946356e-01,\n",
      "              -1.01273581e-01, -1.97538316e-01, -4.21236232e-02,\n",
      "              -2.26530835e-01,  2.57897556e-01, -1.05747268e-01,\n",
      "               3.22286069e-01,  3.51437107e-02, -4.05520707e-01,\n",
      "               2.39359252e-02, -2.20557712e-02, -4.73655127e-02],\n",
      "             [ 2.76853532e-01, -1.44457698e-01, -9.89380926e-02,\n",
      "               3.47104341e-01, -1.61598306e-02, -2.51961946e-01,\n",
      "              -2.43475493e-02, -3.85024846e-01,  9.12255794e-02,\n",
      "              -2.82327265e-01, -1.35130107e-01,  3.42624664e-01,\n",
      "               3.37731928e-01,  2.44988456e-01, -6.11820444e-02],\n",
      "             [ 2.71598816e-01, -1.67596027e-01, -7.99816549e-02,\n",
      "              -1.86424535e-02, -1.24867857e-01,  2.24696562e-01,\n",
      "               2.47132242e-01, -3.84753823e-01,  2.04252243e-01,\n",
      "               4.76379618e-02, -1.70987427e-01,  2.01765746e-01,\n",
      "              -1.58865660e-01, -1.88153550e-01, -1.87683865e-01],\n",
      "             [ 2.67635942e-01,  3.64974409e-01, -4.83548976e-02,\n",
      "               3.17302287e-01, -8.00655037e-03, -8.84973854e-02,\n",
      "               2.39393041e-01, -2.05191821e-01,  2.98161387e-01,\n",
      "               4.12939489e-02,  3.73353273e-01, -1.29213771e-02,\n",
      "               3.88064444e-01, -3.02164376e-01, -1.09527782e-01],\n",
      "             [ 3.93662415e-02, -3.28121066e-01,  5.07963262e-02,\n",
      "              -5.35981469e-02,  2.23266423e-01,  3.09246243e-04,\n",
      "              -2.72782445e-02,  1.72318295e-01, -2.99087435e-01,\n",
      "              -1.50660157e-01,  1.87600795e-02,  3.55571210e-01,\n",
      "               2.06303820e-01,  1.31665822e-02, -3.32961716e-02],\n",
      "             [-1.44158080e-01,  1.32902876e-01,  1.67157948e-01,\n",
      "               2.86985245e-02, -4.96153161e-02, -3.59960973e-01,\n",
      "               3.26851755e-01,  2.51390964e-01, -3.33690852e-01,\n",
      "              -1.42622486e-01, -1.72698915e-01, -2.46725127e-01,\n",
      "              -1.95507295e-02,  2.17494100e-01, -6.87084049e-02],\n",
      "             [ 1.43315673e-01, -5.21935858e-02,  4.27390307e-01,\n",
      "              -1.45552782e-02,  5.14978655e-02,  8.96463320e-02,\n",
      "              -1.20756216e-01, -2.48651445e-01,  2.27677092e-01,\n",
      "               2.17419505e-01, -7.85517767e-02, -4.39446559e-03,\n",
      "              -3.31937559e-02,  3.94727379e-01, -8.97346139e-02],\n",
      "             [-2.47785464e-01,  1.20941050e-01, -1.06237121e-01,\n",
      "              -1.82044581e-02, -1.15215994e-01, -1.08359702e-01,\n",
      "               1.09590150e-01, -1.53881386e-01, -2.36030415e-01,\n",
      "              -7.13651031e-02, -7.16681480e-02,  3.84316534e-01,\n",
      "              -2.50285901e-02,  7.87602514e-02,  1.37455106e-01],\n",
      "             [-9.85348448e-02,  2.59989798e-01, -3.63645524e-01,\n",
      "              -1.65232457e-02, -7.04906583e-02,  1.08477041e-01,\n",
      "               2.39676356e-01, -2.28268102e-01, -2.17864066e-01,\n",
      "               3.64935786e-01, -3.10305774e-01, -3.55145670e-02,\n",
      "               4.67328399e-01, -9.81755555e-02,  1.19568147e-02],\n",
      "             [-8.34548101e-02,  1.14977852e-01,  4.31671962e-02,\n",
      "               2.10428938e-01,  2.69206047e-01, -3.91314067e-02,\n",
      "              -1.65004715e-01, -1.58958912e-01, -6.04332946e-02,\n",
      "               6.16562292e-02,  1.04169454e-02, -1.38164442e-02,\n",
      "               2.25921005e-01,  7.74882883e-02,  6.59611672e-02],\n",
      "             [-3.33960086e-01, -1.34736851e-01,  5.91437370e-02,\n",
      "               7.04332767e-03, -3.97503167e-01,  9.97330472e-02,\n",
      "               2.13668257e-01,  7.68581554e-02, -2.49008581e-01,\n",
      "              -4.36067373e-01, -3.77978534e-01, -2.04937473e-01,\n",
      "               1.70753244e-02,  4.35770482e-01, -1.67447224e-01],\n",
      "             [ 2.75130272e-01,  1.51111811e-01,  1.41167551e-01,\n",
      "               2.29357570e-01, -8.26833248e-02, -4.28926162e-02,\n",
      "              -1.18739866e-01, -4.75106865e-01,  5.28526045e-02,\n",
      "              -3.12575549e-01, -1.16741605e-01,  1.13596722e-01,\n",
      "              -1.29081294e-01,  4.05319929e-01, -2.20293805e-01],\n",
      "             [-4.36499231e-02,  5.19432127e-02,  2.92699873e-01,\n",
      "               2.18307868e-01,  2.31518701e-01, -1.85157970e-01,\n",
      "               7.56429061e-02, -1.30587518e-01, -1.62618354e-01,\n",
      "              -4.69382375e-01, -2.54708320e-01,  4.23803687e-01,\n",
      "              -3.89494896e-01, -2.48679593e-01,  2.57953912e-01],\n",
      "             [-2.97155708e-01,  3.09763342e-01, -1.53159872e-01,\n",
      "               3.00276112e-02,  3.58182460e-01,  2.23884165e-01,\n",
      "               8.91632512e-02, -2.85320310e-03, -9.80710760e-02,\n",
      "               7.23271519e-02,  1.02907894e-02, -3.40505987e-02,\n",
      "              -5.87139316e-02,  1.58778563e-01, -2.28586629e-01],\n",
      "             [ 1.04891457e-01, -3.44674826e-01, -1.03413984e-01,\n",
      "              -1.10771291e-01, -8.49835277e-02,  3.04869805e-02,\n",
      "              -2.40196824e-01, -3.62971127e-01, -4.26439277e-04,\n",
      "              -7.67907724e-02, -5.42494170e-02,  1.87793877e-02,\n",
      "               1.99861735e-01, -4.28949386e-01, -6.95401430e-02],\n",
      "             [ 1.38026133e-01, -7.39240870e-02, -1.35016710e-01,\n",
      "              -1.12478606e-01,  3.07262063e-01,  3.56640756e-01,\n",
      "              -2.65928835e-01, -3.05330724e-01,  2.76131749e-01,\n",
      "              -1.52504919e-02,  3.25810343e-01,  2.58301824e-01,\n",
      "               4.98943664e-02, -1.06995068e-01, -1.16501965e-01],\n",
      "             [-2.65875489e-01, -2.23094210e-01,  2.26308390e-01,\n",
      "              -2.59531289e-01, -3.62237692e-02, -3.84364158e-01,\n",
      "              -5.09146079e-02, -7.14359656e-02,  7.71675706e-02,\n",
      "              -1.93252057e-01, -1.45089224e-01,  1.79051802e-01,\n",
      "              -2.01280296e-01, -8.22212398e-02,  1.78689167e-01],\n",
      "             [-1.30008310e-01, -1.91680267e-02,  3.19532752e-01,\n",
      "              -1.17530532e-01, -2.88273375e-02,  2.34196305e-01,\n",
      "              -2.98909366e-01,  3.32352191e-01, -5.92888296e-02,\n",
      "              -1.09298654e-01, -2.89064292e-02,  2.41845205e-01,\n",
      "              -3.93259227e-01,  3.64573717e-01,  4.45926279e-01],\n",
      "             [ 3.29803601e-02, -1.52929246e-01,  3.31703842e-01,\n",
      "              -4.68752593e-01, -2.20471472e-01, -3.98101006e-03,\n",
      "               4.26751673e-02, -2.23659635e-01, -2.09786639e-01,\n",
      "              -2.05500439e-01,  2.08137840e-01,  1.62023857e-01,\n",
      "               1.81225717e-01,  2.13050582e-02,  1.79024234e-01],\n",
      "             [-1.69833824e-01,  4.09094483e-01, -2.30869040e-01,\n",
      "               2.03401327e-01, -2.56787211e-01, -8.33285153e-02,\n",
      "               2.38466971e-02, -8.80062953e-02, -3.80131900e-01,\n",
      "               1.24576434e-01, -1.76273555e-01,  1.26540095e-01,\n",
      "              -9.20180157e-02, -2.40315989e-01,  2.62186617e-01],\n",
      "             [ 2.28121087e-01, -1.66496187e-01,  1.03084087e-01,\n",
      "               1.66469291e-01,  3.13515335e-01, -1.67697668e-01,\n",
      "              -2.62800664e-01,  7.99010620e-02, -3.76159921e-02,\n",
      "               2.76981643e-03,  1.87145337e-01,  1.98038872e-02,\n",
      "               5.12611866e-02,  3.86081599e-02,  1.45139933e-01],\n",
      "             [-1.02206640e-01, -2.58613616e-01,  1.79068983e-01,\n",
      "              -2.87829101e-01,  3.04776341e-01, -2.43304849e-01,\n",
      "               1.08661816e-01,  2.34098494e-01,  2.59707004e-01,\n",
      "              -1.14874035e-01, -3.33837718e-02,  1.30291596e-01,\n",
      "               2.65145630e-01,  1.79980442e-01,  2.95536779e-02],\n",
      "             [ 1.28452137e-01, -3.19500297e-01, -6.90518618e-02,\n",
      "               1.66528970e-01,  2.12517634e-01, -4.84798886e-02,\n",
      "               1.23467892e-01, -1.95233226e-01,  6.11188561e-02,\n",
      "               4.22154032e-02, -1.26163229e-01,  1.31243810e-01,\n",
      "               3.65756661e-01,  1.79142192e-01, -7.28282034e-02],\n",
      "             [-6.28187135e-02, -1.64800584e-01, -3.61829370e-01,\n",
      "              -9.21417028e-02, -2.52892021e-02,  1.50692333e-02,\n",
      "              -3.69646579e-01,  7.29750171e-02, -1.22602329e-01,\n",
      "               2.40114257e-01,  2.68009603e-02, -9.09864437e-04,\n",
      "              -4.05036323e-02,  1.65354356e-01,  1.75576255e-01]],            dtype=float32), DeviceArray([-0.03376886, -0.00288483,  0.04822405,  0.00964503,\n",
      "             -0.04407851, -0.04513929,  0.02430333,  0.02992898,\n",
      "             -0.02025856, -0.01913235, -0.08598366,  0.00473935,\n",
      "              0.01699474, -0.02820001,  0.01345122], dtype=float32)), (), (DeviceArray([[ 0.23809282, -0.5005227 , -0.10969611,  0.0666615 ,\n",
      "              -0.01058972,  0.43188027,  0.09266079,  0.27215967],\n",
      "             [ 0.36702582,  0.3612049 ,  0.4103505 ,  0.29916817,\n",
      "              -0.32269928,  0.2140048 , -0.27239224,  0.19632143],\n",
      "             [ 0.19586387, -0.24380095, -0.18021023,  0.40611523,\n",
      "               0.16107057, -0.117658  , -0.4696861 ,  0.36912212],\n",
      "             [-0.56228137, -0.00756719,  0.18695138,  0.24438526,\n",
      "               0.01839492, -0.1369225 ,  0.01804907, -0.03050356],\n",
      "             [-0.16516778,  0.19166285, -0.07693551, -0.11446559,\n",
      "              -0.35520276,  0.17312445, -0.07216317,  0.45134246],\n",
      "             [-0.2688479 ,  0.12479939,  0.4795036 ,  0.2788222 ,\n",
      "               0.16599202,  0.47520286,  0.5716269 ,  0.09055096],\n",
      "             [ 0.20154762, -0.00941259,  0.07133609,  0.04694589,\n",
      "               0.1229729 , -0.26983166,  0.32201138,  0.22442925],\n",
      "             [ 0.13150266, -0.17775972,  0.06820236,  0.55790335,\n",
      "              -0.2703583 , -0.43817112,  0.11057483,  0.23614214],\n",
      "             [ 0.06520817,  0.29820043,  0.06941655, -0.5272624 ,\n",
      "              -0.14448637,  0.17394158,  0.6328406 ,  0.4928682 ],\n",
      "             [-0.4077516 , -0.58265126, -0.18208975, -0.25452247,\n",
      "              -0.20577869,  0.24490441,  0.0132821 , -0.05710851],\n",
      "             [-0.25134066,  0.3057414 , -0.10746676, -0.61276627,\n",
      "              -0.0258506 ,  0.399322  , -0.31865647, -0.09370369],\n",
      "             [-0.31341532, -0.05702709, -0.0182786 , -0.5565377 ,\n",
      "               0.22618876, -0.3674944 , -0.10608045, -0.63437784],\n",
      "             [-0.43289706, -0.4159588 ,  0.06029119, -0.6493403 ,\n",
      "              -0.12100775, -0.43746042,  0.11575428,  0.31134212],\n",
      "             [-0.25483575, -0.3998924 , -0.01719618,  0.13531257,\n",
      "               0.01030489, -0.11703007, -0.22380425, -0.51010835],\n",
      "             [ 0.13464046, -0.50561255,  0.20431633, -0.23133531,\n",
      "              -0.5747803 , -0.05846208,  0.48896724, -0.31830662]],            dtype=float32), DeviceArray([ 0.01850034,  0.01343862, -0.04200351,  0.07070687,\n",
      "             -0.00341732, -0.15688053,  0.0073557 ,  0.06442203],            dtype=float32)), (), (DeviceArray([[ 0.2619252 ],\n",
      "             [-0.16194397],\n",
      "             [-0.0270725 ],\n",
      "             [ 0.35862866],\n",
      "             [ 0.06392097],\n",
      "             [-0.8833924 ],\n",
      "             [ 0.20537591],\n",
      "             [ 0.89231026]], dtype=float32), DeviceArray([0.49853265], dtype=float32))]\n",
      "auc=0.7996053178230162\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "\n",
    "# Load the data\n",
    "x1, _ = load_train_dataset(party_id=1)\n",
    "x2, y = load_train_dataset(party_id=2)\n",
    "\n",
    "\n",
    "# Hyperparameter\n",
    "batch_size = 5\n",
    "epochs = 15\n",
    "learning_rate = 0.1\n",
    "\n",
    "\n",
    "# Load the data\n",
    "x1, _ = load_train_dataset(party_id=1)\n",
    "x2, y = load_train_dataset(party_id=2)\n",
    "\n",
    "init_params = init_state(learning_rate)\n",
    "\n",
    "params = train(x1, x2, y, init_params,learning_rate, epochs, batch_size)\n",
    "\n",
    "print(params)\n",
    "\n",
    "X_test, y_test = load_test_dataset()\n",
    "auc = validate_model(params, X_test, y_test)\n",
    "print(f'auc={auc}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may complain AUC here is quite low. But in the end, we are not an AI course. Just keep the number in mind, we are going to repeat the trainning with SPU. Let's do that magic!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Train a Model with SPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-06-16 13:01:25.635121: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "\n",
    "sf.init(['alice', 'bob'], num_cpus=8, log_to_driver=True)\n",
    "\n",
    "alice, bob = sf.PYU('alice'), sf.PYU('bob')\n",
    "spu = sf.SPU(sf.utils.testing.cluster_def(['alice', 'bob']))\n",
    "\n",
    "x1, _ = alice(load_train_dataset)(party_id=1)\n",
    "x2, y = bob(load_train_dataset)(party_id=2)\n",
    "\n",
    "\n",
    "device = spu\n",
    "x1_, x2_, y_ = x1.to(device), x2.to(device), y.to(device)\n",
    "init_params_ = sf.to(spu, lambda: init_state(learning_rate))\n",
    "\n",
    "params_spu = spu(train, static_argnames=['learning_rate', 'epochs', 'batch_size'])(\n",
    "    x1_, x2_, y_,init_params_, learning_rate=learning_rate, epochs=epochs, batch_size=batch_size\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check params from SPU program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2460284)\u001b[0m 2022-06-16 13:01:30.776125: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n",
      "\u001b[2m\u001b[36m(_run pid=2460288)\u001b[0m 2022-06-16 13:01:30.800513: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n",
      "\u001b[2m\u001b[36m(pid=2460290)\u001b[0m 2022-06-16 13:01:30.804987: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n",
      "\u001b[2m\u001b[36m(_run pid=2460286)\u001b[0m 2022-06-16 13:01:30.813806: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n",
      "\u001b[2m\u001b[36m(pid=2460289)\u001b[0m 2022-06-16 13:01:30.857055: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib:/opt/rh/devtoolset-10/root/usr/lib64/dyninst:/opt/rh/devtoolset-10/root/usr/lib/dyninst:/opt/rh/devtoolset-10/root/usr/lib64:/opt/rh/devtoolset-10/root/usr/lib\n",
      "\u001b[2m\u001b[36m(_run pid=2460288)\u001b[0m 2022-06-16 13:01:32,418,418 WARNING [xla_bridge.py:backends:265] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=2460286)\u001b[0m 2022-06-16 13:01:32,508,508 WARNING [xla_bridge.py:backends:265] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=2460284)\u001b[0m 2022-06-16 13:01:32,456,456 WARNING [xla_bridge.py:backends:265] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460290)\u001b[0m I0616 13:01:32.504386 2460290 external/com_github_brpc_brpc/src/brpc/server.cpp:1065] Server[yasl::link::internal::ReceiverServiceImpl] is serving on port=63901.\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460290)\u001b[0m I0616 13:01:32.504480 2460290 external/com_github_brpc_brpc/src/brpc/server.cpp:1068] Check out http://i85c08157.eu95sqa:63901 in web browser.\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460290)\u001b[0m I0616 13:01:32.605220 2461606 external/com_github_brpc_brpc/src/brpc/socket.cpp:2202] Checking Socket{id=0 addr=127.0.0.1:33383} (0x55629fc1ac40)\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460289)\u001b[0m I0616 13:01:32.671928 2460289 external/com_github_brpc_brpc/src/brpc/server.cpp:1065] Server[yasl::link::internal::ReceiverServiceImpl] is serving on port=33383.\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460289)\u001b[0m I0616 13:01:32.672004 2460289 external/com_github_brpc_brpc/src/brpc/server.cpp:1068] Check out http://i85c08157.eu95sqa:33383 in web browser.\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460290)\u001b[0m I0616 13:01:35.605612 2461638 external/com_github_brpc_brpc/src/brpc/socket.cpp:2262] Revived Socket{id=0 addr=127.0.0.1:33383} (0x55629fc1ac40) (Connectable)\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460290)\u001b[0m WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2460289)\u001b[0m WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(array([[ 5.21224737e-02,  2.25653112e-01,  1.42217025e-01,\n",
      "         2.22038493e-01,  1.43407062e-01, -1.17939964e-01,\n",
      "        -2.54890293e-01,  5.02917618e-02,  1.92605436e-01,\n",
      "        -1.66738570e-01,  1.35716438e-01,  9.23201591e-02,\n",
      "        -1.45163283e-01, -2.06950307e-03,  1.02719918e-01,\n",
      "         1.68477565e-01, -1.94680765e-01, -1.68111518e-01,\n",
      "        -2.91620046e-02,  2.47961149e-01, -1.06470272e-01,\n",
      "         1.98481962e-01, -7.39427507e-02, -2.39929244e-01,\n",
      "        -1.13915399e-01,  3.78076851e-01,  1.65145919e-01,\n",
      "         7.91657716e-02, -2.19707474e-01, -7.42163062e-02],\n",
      "       [-2.36553714e-01,  3.90677869e-01, -3.25132608e-02,\n",
      "         1.61681026e-01,  3.65113467e-02, -3.41343939e-01,\n",
      "         3.67962658e-01,  3.51848811e-01, -3.00575048e-02,\n",
      "         1.80104598e-01,  1.01788193e-01,  1.52144164e-01,\n",
      "         2.02847883e-01, -1.55509651e-01, -3.00182939e-01,\n",
      "         2.86474586e-01, -2.53845513e-01, -1.08098716e-01,\n",
      "         2.64337212e-02, -1.04821473e-02, -2.38679454e-01,\n",
      "        -1.13445535e-01, -2.02009663e-01,  3.22062969e-02,\n",
      "        -4.65877205e-02,  1.39607534e-01, -5.21823913e-02,\n",
      "         1.08920917e-01,  1.34518713e-01, -1.11260951e-01],\n",
      "       [-1.82456166e-01,  1.64763957e-01,  6.84951693e-02,\n",
      "         9.42942500e-02, -1.86391711e-01, -8.24576914e-02,\n",
      "        -1.30458355e-01, -4.10286695e-01,  2.35470906e-01,\n",
      "         1.68666959e-01,  4.27838713e-02,  2.72187442e-02,\n",
      "         1.27330750e-01,  1.29255652e-02, -1.90905690e-01,\n",
      "        -4.20718789e-02, -2.68373162e-01,  7.80923069e-02,\n",
      "         3.01362425e-02,  1.55862629e-01, -8.09711963e-02,\n",
      "        -3.08049977e-01, -3.83413434e-02, -3.55214238e-01,\n",
      "        -1.95272222e-01,  4.43269908e-02, -2.24495083e-01,\n",
      "        -1.96313322e-01,  1.26856342e-01, -1.45506412e-01],\n",
      "       [-2.13144869e-01,  1.93932906e-01,  4.28857356e-02,\n",
      "         1.64105698e-01,  9.99248475e-02,  2.14000121e-01,\n",
      "         2.22676069e-01,  6.49158061e-02, -3.63320112e-02,\n",
      "         1.22277334e-01, -6.17514402e-02,  7.48825520e-02,\n",
      "         8.98207873e-02, -1.65753409e-01, -8.66934806e-02,\n",
      "         3.21247995e-01, -1.47644356e-01, -1.28632784e-02,\n",
      "         8.07127357e-02, -2.50974298e-01,  1.47477746e-01,\n",
      "        -4.01914179e-01, -1.30079523e-01, -3.07242274e-01,\n",
      "         1.55475691e-01,  5.10867238e-02, -3.88409317e-01,\n",
      "        -3.00782174e-02, -1.77918017e-01, -3.38865727e-01],\n",
      "       [ 2.84495503e-01, -1.35480687e-01,  1.12002239e-01,\n",
      "        -1.12347975e-01,  2.46275723e-01, -1.25682205e-01,\n",
      "        -1.81171298e-03, -5.22768795e-02, -8.46869349e-02,\n",
      "        -5.46463877e-02, -2.57571638e-01,  2.42745310e-01,\n",
      "        -1.18751675e-01,  2.85120219e-01, -1.30204499e-01,\n",
      "        -6.94652647e-02,  1.56549349e-01,  6.14302903e-02,\n",
      "        -5.26835620e-02,  2.50621140e-03,  4.78822589e-02,\n",
      "         1.76090807e-01, -3.95695120e-02,  1.18063867e-01,\n",
      "        -2.57923543e-01, -3.27984095e-01, -2.13895574e-01,\n",
      "         2.56528765e-01, -3.49479169e-01, -6.18286133e-02],\n",
      "       [-1.76592216e-01,  1.95803612e-01,  2.59060621e-01,\n",
      "         1.27929822e-01,  7.68185854e-02,  4.85150516e-02,\n",
      "        -1.43386230e-01,  1.37525052e-02,  1.24047309e-01,\n",
      "        -7.62728155e-02, -1.63935125e-03, -6.84379935e-02,\n",
      "        -3.60604227e-01,  6.30641282e-02, -1.33007586e-01,\n",
      "         3.05945873e-01,  1.99110359e-01,  3.22232544e-02,\n",
      "         2.40112156e-01, -2.98525453e-01,  2.31414884e-02,\n",
      "         4.47368622e-03, -6.55936152e-02, -1.17711782e-01,\n",
      "         1.33778527e-01,  3.32259685e-01, -3.07792604e-01,\n",
      "         1.24230683e-01, -5.43169081e-02,  4.45076078e-02],\n",
      "       [-2.45652869e-01, -1.72874779e-01, -1.58959270e-01,\n",
      "         1.16324738e-01, -2.48291269e-01, -2.41469592e-01,\n",
      "        -3.95276248e-02,  3.69638979e-01,  6.30413145e-02,\n",
      "        -5.40091097e-02,  2.09700406e-01, -4.00042534e-01,\n",
      "        -1.47648722e-01,  6.29002750e-02, -3.33853662e-01,\n",
      "        -1.53586447e-01, -2.30128273e-01, -1.30878031e-01,\n",
      "         5.11693507e-02, -3.73797715e-02,  1.28014624e-01,\n",
      "         1.09155089e-01, -9.46462303e-02,  9.11422223e-02,\n",
      "        -1.22984871e-01, -3.38232070e-02,  2.96984941e-01,\n",
      "         1.79490834e-01, -3.32677066e-02,  2.96666026e-02],\n",
      "       [-1.30745500e-01, -1.24384150e-01, -3.38499695e-02,\n",
      "         3.48127156e-01,  5.10015637e-02,  9.22064185e-02,\n",
      "         3.15938234e-01, -1.61012575e-01, -8.14388096e-02,\n",
      "        -5.18919230e-02,  1.62081063e-01,  3.27567905e-01,\n",
      "        -1.51928216e-02, -1.68833151e-01, -3.21797788e-01,\n",
      "         3.46298337e-01,  2.55124569e-01,  2.24370107e-01,\n",
      "         2.12284103e-01,  2.16590822e-01,  3.26328129e-02,\n",
      "         2.21246615e-01,  3.73068154e-01,  1.70303017e-01,\n",
      "         2.16773152e-02, -1.61344066e-01,  1.40263379e-01,\n",
      "         4.06199545e-02, -1.64664805e-01,  6.31624311e-02],\n",
      "       [ 7.62210935e-02,  1.80135384e-01,  2.07086876e-01,\n",
      "        -2.31236219e-02, -5.94749302e-02,  3.13537508e-01,\n",
      "        -7.07127005e-02,  5.42959571e-02, -4.78008687e-02,\n",
      "         2.33757809e-01,  3.19292545e-01, -2.42361709e-01,\n",
      "        -1.59331918e-01,  8.17421079e-02, -8.40188116e-02,\n",
      "        -1.77993029e-01,  2.76709110e-01, -2.64198482e-02,\n",
      "         2.40321398e-01, -1.48702189e-01,  1.35500848e-01,\n",
      "         1.54475376e-01,  3.31415117e-01, -7.83364326e-02,\n",
      "        -1.96747780e-02,  1.50718108e-01,  2.75711417e-02,\n",
      "         2.92210579e-01,  1.25457793e-01, -1.37014270e-01],\n",
      "       [-1.54811621e-01,  2.58745551e-01, -3.99873674e-01,\n",
      "         1.32840917e-01, -5.03840148e-02, -1.33610189e-01,\n",
      "        -3.32632422e-01,  5.91315925e-02, -1.07699126e-01,\n",
      "        -8.45145881e-02, -1.09385028e-01,  7.60867447e-02,\n",
      "        -1.28735870e-01,  2.23394170e-01,  1.21435210e-01,\n",
      "        -2.82440633e-02, -2.88350880e-01,  6.86985552e-02,\n",
      "        -7.02297837e-02, -1.59900114e-01, -1.61253199e-01,\n",
      "        -1.46421507e-01,  6.38383925e-02,  1.64755732e-02,\n",
      "         1.49303243e-01,  4.48648781e-02, -5.04779816e-03,\n",
      "        -3.15861106e-01, -4.52296436e-02, -5.61970472e-03],\n",
      "       [ 2.41484672e-01,  7.55192041e-02,  2.95221388e-01,\n",
      "         1.46864414e-01,  9.80748385e-02,  2.25508973e-01,\n",
      "        -6.05033338e-02, -1.74552232e-01,  3.88410538e-02,\n",
      "         2.18085364e-01, -9.08160210e-02,  8.86506587e-02,\n",
      "         1.37923896e-01,  4.19318944e-01,  7.50737339e-02,\n",
      "        -2.17468202e-01,  1.65094569e-01, -1.65486336e-03,\n",
      "        -9.91209149e-02,  1.45306885e-01, -1.13687873e-01,\n",
      "         5.27484268e-02,  1.29981056e-01,  3.83183241e-01,\n",
      "        -3.15553039e-01,  1.89110160e-01, -1.34139583e-01,\n",
      "         1.36779934e-01,  2.81374305e-02, -1.72807977e-01],\n",
      "       [-2.15493441e-02,  4.04925644e-02, -8.68780166e-02,\n",
      "         6.79320842e-02,  9.82559323e-02,  6.61205947e-02,\n",
      "        -2.02870637e-01, -1.17738158e-01,  1.63120881e-01,\n",
      "         2.78532803e-01, -3.04895878e-01, -2.78360248e-02,\n",
      "        -1.67953745e-01,  3.51908863e-01, -5.97648919e-02,\n",
      "         1.47835299e-01, -2.78262317e-01, -8.54300708e-02,\n",
      "        -6.35778904e-03, -1.14808097e-01,  5.83171844e-03,\n",
      "         9.01198089e-02, -1.17785096e-01,  1.15210667e-01,\n",
      "        -1.71339333e-01,  1.55093044e-01, -3.24716568e-02,\n",
      "        -3.46243024e-01,  6.52319789e-02,  1.80368528e-01],\n",
      "       [ 7.85570890e-02,  2.71920562e-02, -1.24795094e-01,\n",
      "         2.27356747e-01,  2.99579352e-02,  2.80274004e-02,\n",
      "        -2.30154023e-01,  2.21238405e-01, -3.20344716e-02,\n",
      "        -1.70208976e-01, -2.27579504e-01, -2.73940355e-01,\n",
      "        -1.20936826e-01,  1.36115164e-01, -1.06474981e-01,\n",
      "        -6.11828566e-02, -5.84995747e-02, -2.57405996e-01,\n",
      "         3.59185845e-01, -3.75315994e-02,  4.23692018e-02,\n",
      "         7.34555274e-02, -3.62589955e-01, -4.30174768e-02,\n",
      "         1.28380165e-01,  3.00772786e-01, -1.82752147e-01,\n",
      "        -1.48194656e-01, -3.67040575e-01, -3.29672158e-01],\n",
      "       [ 2.83481181e-03,  1.81894600e-02, -2.20816538e-01,\n",
      "        -1.68947875e-03, -1.24119565e-01,  1.66682243e-01,\n",
      "         1.04996771e-01,  2.45537356e-01,  3.30403775e-01,\n",
      "        -3.11017454e-01,  3.88675988e-01, -2.32338458e-01,\n",
      "         5.28738350e-02, -4.66596335e-02,  6.83461577e-02,\n",
      "         3.46443444e-01, -2.03109890e-01,  3.64813238e-01,\n",
      "        -2.89645880e-01,  1.82707608e-02,  3.27082276e-01,\n",
      "         2.45703429e-01,  4.19372767e-02,  3.54042053e-02,\n",
      "        -2.76969671e-02, -2.81784892e-01,  2.04562888e-01,\n",
      "        -3.01399291e-01,  3.56618315e-02,  4.00729477e-01],\n",
      "       [ 2.89673388e-01,  1.94585279e-01, -7.54357874e-02,\n",
      "        -1.89162433e-01,  9.51814502e-02, -1.66590184e-01,\n",
      "         2.11697072e-02, -9.17358100e-02, -2.53419816e-01,\n",
      "        -1.87644809e-02, -1.17323250e-02,  4.42565233e-02,\n",
      "         1.09574080e-01,  3.88969958e-01, -4.52404916e-02,\n",
      "        -4.87569422e-02, -3.13752860e-01, -1.72892362e-02,\n",
      "         1.44217759e-02, -2.87221372e-03, -2.70475864e-01,\n",
      "        -2.28503153e-01, -9.87705290e-02, -6.37385249e-02,\n",
      "        -1.52837306e-01,  2.69068778e-02, -3.35106552e-02,\n",
      "         1.30309016e-02, -9.01488960e-03,  4.21783924e-02],\n",
      "       [ 3.43039930e-02,  1.50743067e-01,  8.29366148e-02,\n",
      "        -1.77595943e-01,  2.46896744e-02, -3.33037704e-01,\n",
      "        -9.81211662e-04,  1.99418977e-01,  3.33111644e-01,\n",
      "        -1.02562189e-01, -3.86553824e-01,  2.62529999e-02,\n",
      "        -2.49244168e-01,  5.33246398e-02, -5.41717112e-02,\n",
      "        -2.58140266e-01, -1.37030661e-01,  1.61672294e-01,\n",
      "         1.15453362e-01, -3.02354425e-01, -8.03447217e-02,\n",
      "        -1.95433080e-01,  3.14462364e-01,  2.45914832e-01,\n",
      "        -7.07805753e-02,  4.00525600e-01,  1.04964003e-01,\n",
      "         1.60914063e-02, -1.69112235e-01,  3.16944420e-02],\n",
      "       [-8.87536258e-02,  1.58608600e-01, -2.44522899e-01,\n",
      "         2.79460996e-02, -2.89275110e-01,  3.18608642e-01,\n",
      "        -1.31275356e-02,  1.51372477e-01, -1.29635155e-01,\n",
      "        -1.82447344e-01, -1.96553543e-01,  9.16007906e-02,\n",
      "        -1.75659880e-01,  1.69408962e-01,  4.26574051e-02,\n",
      "        -3.11639190e-01,  2.54636467e-01,  1.29695728e-01,\n",
      "        -6.24269843e-02,  2.07581341e-01, -1.62594169e-02,\n",
      "        -8.78837854e-02,  2.46837571e-01, -1.96480006e-02,\n",
      "         1.33252591e-02,  8.69227946e-02,  1.78894758e-01,\n",
      "         2.63514340e-01, -1.26071960e-01, -2.81487763e-01],\n",
      "       [-2.72540390e-01,  3.90539885e-01,  1.52694240e-01,\n",
      "         3.31544876e-03, -1.77479222e-01, -2.27627307e-01,\n",
      "        -2.12363526e-01, -1.39447004e-02,  2.46945098e-01,\n",
      "         1.44541472e-01, -2.77292937e-01,  2.96193779e-01,\n",
      "         4.70548123e-02, -8.88011903e-02, -9.71731246e-02,\n",
      "        -2.62320638e-02, -2.95179188e-02,  1.86319903e-01,\n",
      "         2.01716110e-01,  1.29090026e-01, -1.41982451e-01,\n",
      "        -2.14984268e-01,  9.17072296e-02, -4.58690673e-02,\n",
      "        -2.05015630e-01, -2.43071705e-01,  1.10963732e-02,\n",
      "         7.75285661e-02,  1.24029964e-02,  7.20824301e-02],\n",
      "       [-8.60645324e-02,  7.43355155e-02, -8.76725316e-02,\n",
      "        -1.68098658e-02, -3.40442061e-01,  3.35628390e-02,\n",
      "        -3.70488763e-02,  3.24337959e-01, -5.66657782e-02,\n",
      "        -2.27891460e-01, -1.25499204e-01,  2.12699145e-01,\n",
      "        -1.46616101e-02,  1.23815641e-01, -4.73732054e-02,\n",
      "        -1.34847909e-02,  2.06255853e-01, -5.21928072e-03,\n",
      "        -1.25348747e-01,  6.38622791e-02,  2.05208302e-01,\n",
      "         1.96026221e-01,  4.12296832e-01,  2.49014318e-01,\n",
      "        -2.93397307e-01, -2.37924501e-01,  1.07741505e-02,\n",
      "         2.44041055e-01,  6.08813614e-02,  2.70513594e-02],\n",
      "       [ 1.54369712e-01, -6.89209700e-02,  2.13278234e-02,\n",
      "        -1.44704819e-01,  7.00449198e-02,  3.90132070e-02,\n",
      "        -1.51024625e-01,  3.11700165e-01,  1.62612468e-01,\n",
      "         7.81369060e-02,  8.33999217e-02,  1.94705576e-01,\n",
      "        -3.38721573e-02,  3.80493701e-03, -7.93460011e-03,\n",
      "        -3.71055782e-01,  1.88404590e-01,  1.32495701e-01,\n",
      "        -2.16117650e-02, -9.45566148e-02, -1.32949337e-01,\n",
      "        -2.76982188e-02,  2.63534904e-01, -1.80030167e-01,\n",
      "         2.15611666e-01, -8.54085386e-02,  6.10715449e-02,\n",
      "        -4.57692742e-02,  1.37957945e-01, -1.13301516e-01],\n",
      "       [-1.10474154e-01, -2.33840957e-01, -1.24218017e-01,\n",
      "         6.95376694e-02,  1.86662674e-02, -1.75065234e-01,\n",
      "         2.81300247e-02,  1.19689360e-01, -2.70128489e-01,\n",
      "         1.95977256e-01,  1.92392498e-01,  2.52708197e-01,\n",
      "        -1.32820055e-01,  6.47971779e-02, -9.03438926e-02,\n",
      "         7.40513057e-02,  6.70711696e-03, -3.97984982e-02,\n",
      "         1.77754119e-01, -2.30836630e-01,  3.55218112e-01,\n",
      "         2.82210827e-01, -1.78824499e-01,  6.64580166e-02,\n",
      "         1.27407983e-01,  1.16543323e-01,  1.38543531e-01,\n",
      "         1.29148975e-01,  3.33635449e-01,  1.66850686e-01],\n",
      "       [ 1.98070556e-01,  8.14639181e-02, -9.22388583e-02,\n",
      "        -1.13172159e-01, -5.22312373e-02,  8.12314451e-02,\n",
      "        -3.92124206e-01, -1.01250544e-01,  2.19034493e-01,\n",
      "        -1.70103595e-01,  2.79538751e-01, -1.13394618e-01,\n",
      "        -1.51904285e-01, -1.66047215e-02,  1.89192876e-01,\n",
      "        -2.87384778e-01,  2.03953311e-01,  1.80972815e-01,\n",
      "         1.43214494e-01,  1.88757628e-01, -1.74619913e-01,\n",
      "         3.13838303e-01, -6.46850318e-02, -2.46058464e-01,\n",
      "        -1.88325718e-01,  7.07657188e-02,  7.73985386e-02,\n",
      "         5.96041679e-02, -1.85137302e-01,  5.19104600e-02],\n",
      "       [ 3.71416450e-01, -2.90058851e-02, -1.86454117e-01,\n",
      "        -1.15898252e-02, -1.55883491e-01, -1.18888110e-01,\n",
      "         4.19633240e-02,  4.30098236e-01,  5.93752712e-02,\n",
      "        -1.17486268e-01,  1.49890855e-01,  1.61925822e-01,\n",
      "         5.26798218e-02,  2.31389567e-01, -1.91654652e-01,\n",
      "         9.47254896e-02, -1.88520849e-02,  2.34299496e-01,\n",
      "         2.27973551e-01,  1.80854887e-01,  1.43894553e-03,\n",
      "        -1.05802253e-01, -1.18630022e-01,  1.46356925e-01,\n",
      "        -4.84475344e-02,  5.88763505e-02, -3.32957685e-01,\n",
      "         5.97254932e-02,  2.21563116e-01,  1.50314972e-01],\n",
      "       [-1.35240108e-01, -1.21922791e-01, -1.33650333e-01,\n",
      "        -1.02593690e-01, -1.78822562e-01, -3.92985553e-01,\n",
      "         3.86028588e-01,  1.58652663e-04, -6.60253465e-02,\n",
      "        -2.11674392e-01, -3.62494141e-02, -1.09350652e-01,\n",
      "         2.18345299e-01, -7.22628236e-02,  6.00301176e-02,\n",
      "        -1.98951066e-02,  3.85129452e-02,  2.57069916e-01,\n",
      "         8.68800879e-02,  1.24661595e-01,  3.28119695e-02,\n",
      "         6.96161389e-03, -1.00336775e-01, -3.21803957e-01,\n",
      "        -6.43377453e-02, -3.76082510e-02, -1.00084111e-01,\n",
      "         3.22159410e-01,  1.86191455e-01, -2.68863738e-01],\n",
      "       [-7.22350925e-02,  9.20040309e-02, -3.33778083e-01,\n",
      "        -3.25040460e-01, -2.19114453e-01,  1.69675574e-01,\n",
      "         2.28981987e-01, -1.12164930e-01,  8.49121809e-03,\n",
      "        -2.36354381e-01,  2.44443297e-01, -1.94931805e-01,\n",
      "         9.91718769e-02,  8.44576061e-02, -1.72666132e-01,\n",
      "         4.42800820e-02, -1.23017579e-01,  3.32668424e-02,\n",
      "        -3.07134926e-01,  1.45027772e-01, -1.72679961e-01,\n",
      "        -1.28169686e-01, -2.04624429e-01,  8.39755386e-02,\n",
      "         3.24039012e-01,  1.49567783e-01, -1.39526129e-02,\n",
      "         1.65045902e-01,  2.51272738e-01,  5.56008667e-02],\n",
      "       [ 7.67881423e-02, -1.54560968e-01, -1.35633782e-01,\n",
      "        -3.76705378e-02,  2.88041711e-01, -2.13424742e-01,\n",
      "         3.19540560e-01,  2.10555166e-01,  1.52228758e-01,\n",
      "        -2.10315466e-01,  1.11538559e-01,  1.12477750e-01,\n",
      "         2.68751323e-01,  1.75382808e-01,  4.28125262e-03,\n",
      "        -2.38750130e-02, -1.72972217e-01, -6.97739720e-02,\n",
      "         5.90839982e-02, -3.05287719e-01,  2.03098953e-02,\n",
      "        -9.70764458e-02, -1.03239924e-01, -2.27292329e-02,\n",
      "         2.63447285e-01,  3.56167257e-02, -5.88473976e-02,\n",
      "        -1.98080823e-01,  2.52847522e-01,  6.78236485e-02],\n",
      "       [ 3.62003863e-01, -3.70889008e-02,  2.65850663e-01,\n",
      "         6.54557496e-02,  1.52984396e-01, -2.51893699e-02,\n",
      "         2.87941694e-02, -1.38501033e-01,  1.56723320e-01,\n",
      "        -2.43221134e-01, -1.85205787e-01,  7.18113780e-02,\n",
      "         1.65021122e-01, -1.05036780e-01, -1.58209205e-01,\n",
      "        -2.15166718e-01,  8.02514255e-02, -2.78826296e-01,\n",
      "        -2.72632986e-02, -5.86647242e-02, -9.58307087e-03,\n",
      "         2.69430399e-01,  3.70774209e-01,  1.51875660e-01,\n",
      "        -2.08369404e-01, -6.35239631e-02,  1.64800867e-01,\n",
      "        -1.59679800e-01, -9.01359022e-02, -3.13538253e-01],\n",
      "       [-2.66990900e-01, -2.93253720e-01, -3.75201762e-01,\n",
      "         5.79940975e-02, -5.38680851e-02,  5.15445769e-02,\n",
      "        -1.94441989e-01,  2.77562499e-01,  1.06624097e-01,\n",
      "         2.95093060e-02,  1.38826609e-01,  1.55660585e-01,\n",
      "         9.79765952e-02, -2.69859195e-01,  2.11702570e-01,\n",
      "        -1.53250426e-01,  2.84165293e-01,  2.54740119e-01,\n",
      "         3.38875949e-02, -1.93462953e-01, -1.74930841e-01,\n",
      "        -8.87168199e-02,  1.60166219e-01, -9.35309976e-02,\n",
      "         2.01478660e-01, -1.91914469e-01, -1.05857968e-01,\n",
      "         9.67087448e-02,  5.61915785e-02,  3.74115556e-02],\n",
      "       [-4.80836034e-02, -7.00737238e-02,  3.11958402e-01,\n",
      "        -2.43130475e-01,  3.71656090e-01, -3.15267324e-01,\n",
      "         1.72800496e-01, -2.13124231e-01, -3.10963660e-01,\n",
      "         1.96563959e-01,  2.35389039e-01, -1.65749654e-01,\n",
      "         9.80474055e-03, -1.29300311e-01, -7.46937245e-02,\n",
      "         1.18998811e-01, -2.02566192e-01,  1.39120445e-01,\n",
      "         1.54781491e-01, -3.16816568e-03,  8.13446790e-02,\n",
      "         1.81269795e-02,  2.81595021e-01,  1.31729126e-01,\n",
      "        -2.29917854e-01,  3.45019400e-01,  2.67447650e-01,\n",
      "         2.59035528e-01,  1.55651569e-01, -2.56225467e-03],\n",
      "       [ 6.15930557e-03,  6.07086122e-02,  7.43321627e-02,\n",
      "        -2.29558349e-02, -1.29004270e-01, -1.20779097e-01,\n",
      "        -1.12963110e-01, -3.19789112e-01,  1.36005774e-01,\n",
      "        -2.26353198e-01,  8.71904343e-02,  3.82686764e-01,\n",
      "         2.32243821e-01,  1.74642175e-01,  4.04576957e-02,\n",
      "        -1.08322740e-01, -1.50881022e-01, -2.06239164e-01,\n",
      "        -2.60582983e-01,  1.00144118e-01,  2.21371815e-01,\n",
      "         2.46227413e-01, -2.86235332e-01,  1.58927217e-01,\n",
      "         8.62326622e-02,  1.77884251e-01, -1.82513297e-01,\n",
      "         5.37948310e-02, -2.14487314e-03,  1.91295892e-02]], dtype=float32), array([ 0.01731433,  0.01454931,  0.01885408,  0.01764621, -0.00197086,\n",
      "       -0.02482608, -0.00320765,  0.02726007, -0.02903186, -0.02995506,\n",
      "       -0.01840489,  0.00923604, -0.03126048,  0.02295962,  0.01063779,\n",
      "       -0.04037455,  0.02124014,  0.0038335 ,  0.00244744, -0.00779128,\n",
      "       -0.02259418, -0.02033043,  0.03011787, -0.0212658 , -0.01007034,\n",
      "       -0.00041412, -0.03257667,  0.01641177,  0.00886433, -0.00069462],\n",
      "      dtype=float32)), (), (array([[ 2.39794090e-01,  1.11328945e-01, -6.55358434e-02,\n",
      "        -1.53777465e-01, -1.58345401e-02, -2.14052558e-01,\n",
      "         4.12578464e-01, -1.79303110e-01, -9.40021276e-02,\n",
      "        -2.57758021e-01, -2.78356820e-02, -2.03546733e-01,\n",
      "        -1.38334900e-01,  2.14947328e-01,  1.98414788e-01],\n",
      "       [-2.03269213e-01,  2.72569358e-02,  2.13408619e-02,\n",
      "         1.91221580e-01, -9.27539319e-02, -1.01883575e-01,\n",
      "        -9.96604562e-03, -1.39726982e-01,  2.97648340e-01,\n",
      "        -5.27602285e-02, -1.32996574e-01, -1.73702270e-01,\n",
      "        -2.69065350e-02,  3.85854483e-01,  3.27751219e-01],\n",
      "       [-2.27373242e-02, -2.26361960e-01, -3.55648100e-02,\n",
      "        -2.46449292e-01, -9.35786068e-02, -6.52750880e-02,\n",
      "         3.50861251e-03, -3.24443161e-01,  3.30272317e-02,\n",
      "         6.77857399e-02, -1.99820548e-02, -1.61494687e-01,\n",
      "         3.19004476e-01, -1.98884234e-01,  5.25489748e-02],\n",
      "       [ 1.89079657e-01, -2.82467723e-01, -2.34362006e-01,\n",
      "         2.82163858e-01, -2.85350740e-01, -8.80528390e-02,\n",
      "        -4.46510136e-01,  2.65960753e-01, -1.03806034e-01,\n",
      "         1.94419786e-01,  1.69521138e-01, -8.81101638e-02,\n",
      "        -1.13381609e-01, -8.65094811e-02, -4.72609699e-02],\n",
      "       [-4.33977842e-01,  4.05445695e-04, -2.06277847e-01,\n",
      "        -3.43710333e-02, -9.84994322e-02, -1.30355805e-02,\n",
      "        -1.78069472e-01,  1.38594538e-01,  2.01298296e-02,\n",
      "        -6.78939223e-02, -1.91350564e-01,  1.01637051e-01,\n",
      "        -1.84102431e-01,  4.55374122e-02,  1.89003497e-02],\n",
      "       [ 3.87100130e-02,  2.38577217e-01,  7.68883228e-02,\n",
      "        -4.95792478e-02,  3.21685791e-01,  2.42307052e-01,\n",
      "         1.17416337e-01,  3.30441415e-01,  2.49463022e-02,\n",
      "        -2.39285782e-01,  1.80279255e-01, -2.65259922e-01,\n",
      "        -1.70731157e-01,  1.92595646e-01,  7.27543235e-02],\n",
      "       [-1.40079170e-01,  4.77774441e-03,  2.86962032e-01,\n",
      "        -1.01265386e-01, -1.97653532e-01, -4.21464145e-02,\n",
      "        -2.25989535e-01,  2.57803589e-01, -1.05742857e-01,\n",
      "         3.22332531e-01,  3.51299644e-02, -4.05459046e-01,\n",
      "         2.42620409e-02, -2.23983973e-02, -4.74648923e-02],\n",
      "       [ 2.76862562e-01, -1.44462705e-01, -9.89008546e-02,\n",
      "         3.47111881e-01, -1.61352009e-02, -2.51957357e-01,\n",
      "        -2.43358612e-02, -3.84998322e-01,  9.12416428e-02,\n",
      "        -2.82317579e-01, -1.35125071e-01,  3.42629224e-01,\n",
      "         3.37768435e-01,  2.44969919e-01, -6.12225384e-02],\n",
      "       [ 2.71600604e-01, -1.67571962e-01, -8.01121294e-02,\n",
      "        -1.86335593e-02, -1.25043750e-01,  2.24736512e-01,\n",
      "         2.47157946e-01, -3.84748429e-01,  2.04266682e-01,\n",
      "         4.76682633e-02, -1.70982346e-01,  2.01754987e-01,\n",
      "        -1.58929169e-01, -1.88026875e-01, -1.87662587e-01],\n",
      "       [ 2.67625362e-01,  3.64946157e-01, -4.84992415e-02,\n",
      "         3.17310482e-01, -7.98653066e-03, -8.84386152e-02,\n",
      "         2.39333451e-01, -2.05214322e-01,  2.98170567e-01,\n",
      "         4.13188636e-02,  3.73357594e-01, -1.28719807e-02,\n",
      "         3.88075948e-01, -3.02084982e-01, -1.09461680e-01],\n",
      "       [ 3.95832807e-02, -3.28108132e-01,  5.09423167e-02,\n",
      "        -5.35945743e-02,  2.23526061e-01,  3.01167369e-04,\n",
      "        -2.71179974e-02,  1.72336340e-01, -2.99085140e-01,\n",
      "        -1.50664732e-01,  1.87656283e-02,  3.55639219e-01,\n",
      "         2.06464499e-01,  1.29646510e-02, -3.35267186e-02],\n",
      "       [-1.43441364e-01,  1.33213788e-01,  1.67769447e-01,\n",
      "         2.87075639e-02, -4.93597537e-02, -3.59961569e-01,\n",
      "         3.27408731e-01,  2.51425892e-01, -3.33686560e-01,\n",
      "        -1.42691106e-01, -1.72692671e-01, -2.46359885e-01,\n",
      "        -1.89049989e-02,  2.16630951e-01, -6.94303215e-02],\n",
      "       [ 1.44663900e-01, -5.21379262e-02,  4.28086877e-01,\n",
      "        -1.45526975e-02,  5.13807386e-02,  8.94423127e-02,\n",
      "        -1.19835809e-01, -2.48696879e-01,  2.27682948e-01,\n",
      "         2.17367962e-01, -7.84611404e-02, -4.10269201e-03,\n",
      "        -3.23567539e-02,  3.93726528e-01, -8.96166265e-02],\n",
      "       [-2.47777447e-01,  1.20911151e-01, -1.06197476e-01,\n",
      "        -1.82008147e-02, -1.14703462e-01, -1.08131856e-01,\n",
      "         1.09589860e-01, -1.53841630e-01, -2.36023694e-01,\n",
      "        -7.13508129e-02, -7.16695786e-02,  3.84531379e-01,\n",
      "        -2.50048637e-02,  7.87962824e-02,  1.37469441e-01],\n",
      "       [-9.76380706e-02,  2.59810925e-01, -3.63053620e-01,\n",
      "        -1.65193975e-02, -7.04343021e-02,  1.08283147e-01,\n",
      "         2.40159720e-01, -2.28284732e-01, -2.17856184e-01,\n",
      "         3.64844322e-01, -3.10198784e-01, -3.51060182e-02,\n",
      "         4.68011916e-01, -9.90075469e-02,  1.22742653e-02],\n",
      "       [-8.34562182e-02,  1.14999309e-01,  4.28011864e-02,\n",
      "         2.10435852e-01,  2.69063473e-01, -3.90328169e-02,\n",
      "        -1.64978698e-01, -1.58941820e-01, -6.04219586e-02,\n",
      "         6.16926253e-02,  1.04072392e-02, -1.38352960e-02,\n",
      "         2.25787848e-01,  7.77026564e-02,  6.60229772e-02],\n",
      "       [-3.33959401e-01, -1.34706646e-01,  5.89306355e-02,\n",
      "         7.04857707e-03, -3.97473037e-01,  9.97875184e-02,\n",
      "         2.13663340e-01,  7.68283308e-02, -2.49001414e-01,\n",
      "        -4.36027348e-01, -3.77961129e-01, -2.04589844e-01,\n",
      "         1.71852261e-02,  4.35841978e-01, -1.67440698e-01],\n",
      "       [ 2.75114417e-01,  1.51113451e-01,  1.41153112e-01,\n",
      "         2.29369670e-01, -8.25335234e-02, -4.28930372e-02,\n",
      "        -1.18737429e-01, -4.75101471e-01,  5.28638214e-02,\n",
      "        -3.12567502e-01, -1.16736934e-01,  1.13632202e-01,\n",
      "        -1.29112586e-01,  4.05371010e-01, -2.20292136e-01],\n",
      "       [-4.36410010e-02,  5.19461334e-02,  2.92714834e-01,\n",
      "         2.18322590e-01,  2.31543735e-01, -1.85120702e-01,\n",
      "         7.56793022e-02, -1.30581975e-01, -1.62602231e-01,\n",
      "        -4.69376862e-01, -2.54704654e-01,  4.23734546e-01,\n",
      "        -3.89467537e-01, -2.48743758e-01,  2.57949114e-01],\n",
      "       [-2.97238588e-01,  3.09765339e-01, -1.53431550e-01,\n",
      "         3.00379843e-02,  3.57985854e-01,  2.23955065e-01,\n",
      "         8.90776366e-02, -2.79456377e-03, -9.80617106e-02,\n",
      "         7.24013299e-02,  1.02955997e-02, -3.39252949e-02,\n",
      "        -5.88560104e-02,  1.59226134e-01, -2.28242025e-01],\n",
      "       [ 1.04995176e-01, -3.44647706e-01, -1.03543982e-01,\n",
      "        -1.10752955e-01, -8.50481838e-02,  3.04953307e-02,\n",
      "        -2.40112007e-01, -3.62978339e-01, -4.27335501e-04,\n",
      "        -7.67635256e-02, -5.42325526e-02,  1.87263638e-02,\n",
      "         1.99837103e-01, -4.28885937e-01, -6.93002194e-02],\n",
      "       [ 1.38648748e-01, -7.38954097e-02, -1.34206176e-01,\n",
      "        -1.12471014e-01,  3.07312131e-01,  3.56791705e-01,\n",
      "        -2.65457153e-01, -3.05273533e-01,  2.76138633e-01,\n",
      "        -1.53627098e-02,  3.25792193e-01,  2.58237422e-01,\n",
      "         5.05699813e-02, -1.08101323e-01, -1.17197156e-01],\n",
      "       [-2.65867561e-01, -2.23087251e-01,  2.26275027e-01,\n",
      "        -2.59516299e-01, -3.62234265e-02, -3.84332299e-01,\n",
      "        -5.09195030e-02, -7.14417994e-02,  7.71806985e-02,\n",
      "        -1.93237886e-01, -1.45074248e-01,  1.79460645e-01,\n",
      "        -2.01161265e-01, -8.22328776e-02,  1.78644389e-01],\n",
      "       [-1.29698217e-01, -1.91744715e-02,  3.19914371e-01,\n",
      "        -1.17517233e-01, -2.88056582e-02,  2.34326616e-01,\n",
      "        -2.98719525e-01,  3.32405657e-01, -5.92833608e-02,\n",
      "        -1.09351203e-01, -2.89273411e-02,  2.41676107e-01,\n",
      "        -3.92921627e-01,  3.64018917e-01,  4.45611060e-01],\n",
      "       [ 3.39959115e-02, -1.53093725e-01,  3.32459927e-01,\n",
      "        -4.68747616e-01, -2.20022246e-01, -4.30546701e-03,\n",
      "         4.31732237e-02, -2.23696575e-01, -2.09781110e-01,\n",
      "        -2.05605090e-01,  2.08349735e-01,  1.62514359e-01,\n",
      "         1.82057247e-01,  2.03865319e-02,  1.78974211e-01],\n",
      "       [-1.69827461e-01,  4.09096479e-01, -2.30982527e-01,\n",
      "         2.03410625e-01, -2.56908804e-01, -8.32471102e-02,\n",
      "         2.38855928e-02, -8.80017281e-02, -3.80121708e-01,\n",
      "         1.24601007e-01, -1.76273659e-01,  1.26525149e-01,\n",
      "        -9.20653492e-02, -2.40216330e-01,  2.62216806e-01],\n",
      "       [ 2.28171796e-01, -1.66493803e-01,  1.03221551e-01,\n",
      "         1.66482970e-01,  3.13554198e-01, -1.67614549e-01,\n",
      "        -2.62774259e-01,  8.00031871e-02, -3.76106203e-02,\n",
      "         2.76772678e-03,  1.87133908e-01,  1.98301822e-02,\n",
      "         5.14633209e-02,  3.84298861e-02,  1.44911990e-01],\n",
      "       [-1.02200523e-01, -2.58575499e-01,  1.79089859e-01,\n",
      "        -2.87817210e-01,  3.04817766e-01, -2.43289828e-01,\n",
      "         1.08641580e-01,  2.34156087e-01,  2.59714335e-01,\n",
      "        -1.14846990e-01, -3.33968699e-02,  1.30755290e-01,\n",
      "         2.65262783e-01,  1.80035874e-01,  2.94272602e-02],\n",
      "       [ 1.28671095e-01, -3.19525421e-01, -6.88779056e-02,\n",
      "         1.66534439e-01,  2.12999523e-01, -4.85396683e-02,\n",
      "         1.23570099e-01, -1.95227697e-01,  6.11242205e-02,\n",
      "         4.22084183e-02, -1.26157284e-01,  1.31496325e-01,\n",
      "         3.65891427e-01,  1.78969011e-01, -7.29027390e-02],\n",
      "       [-6.28064871e-02, -1.64812818e-01, -3.61909330e-01,\n",
      "        -9.21332687e-02, -2.51936466e-02,  1.51226223e-02,\n",
      "        -3.69720221e-01,  7.29784220e-02, -1.22595370e-01,\n",
      "         2.40122572e-01,  2.67935991e-02, -8.51482153e-04,\n",
      "        -4.04837728e-02,  1.65353596e-01,  1.75601244e-01]], dtype=float32), array([-0.03335796, -0.00299449,  0.04807   ,  0.00967184, -0.04413249,\n",
      "       -0.0450432 ,  0.02446948,  0.02991146, -0.020246  , -0.01907897,\n",
      "       -0.0858811 ,  0.00561506,  0.01732561, -0.02812673,  0.0137278 ],\n",
      "      dtype=float32)), (), (array([[ 0.23808274, -0.50051725, -0.10969022,  0.06664677, -0.01056869,\n",
      "         0.43188578,  0.09266037,  0.2720422 ],\n",
      "       [ 0.366984  ,  0.36121023,  0.4103575 ,  0.29919928, -0.32269377,\n",
      "         0.21401024, -0.27237856,  0.1960319 ],\n",
      "       [ 0.19575597, -0.24379551, -0.18022081,  0.40605476,  0.16112107,\n",
      "        -0.11765367, -0.46969986,  0.37025326],\n",
      "       [-0.5622789 , -0.00756182,  0.18691337,  0.24441676,  0.01844142,\n",
      "        -0.13691702,  0.01817606, -0.03054906],\n",
      "       [-0.16521893,  0.1916684 , -0.07693869, -0.11437152, -0.35518903,\n",
      "         0.17312975, -0.07215329,  0.45142734],\n",
      "       [-0.2688534 ,  0.12480481,  0.4795176 ,  0.278813  ,  0.1659974 ,\n",
      "         0.47520757,  0.571603  ,  0.09030092],\n",
      "       [ 0.20145267, -0.00940721,  0.07133275,  0.04697774,  0.12301297,\n",
      "        -0.2698263 ,  0.32198858,  0.2248482 ],\n",
      "       [ 0.13138188, -0.17775431,  0.06820518,  0.5579902 , -0.2703543 ,\n",
      "        -0.43816575,  0.11057664,  0.23617218],\n",
      "       [ 0.06521353,  0.29820585,  0.0694216 , -0.5272582 , -0.14447974,\n",
      "         0.173947  ,  0.6328493 ,  0.4928742 ],\n",
      "       [-0.40775123, -0.5826458 , -0.18208268, -0.25455818, -0.20577785,\n",
      "         0.24490973,  0.01328488, -0.0576129 ],\n",
      "       [-0.25134367,  0.3057468 , -0.1074589 , -0.61274683, -0.02584518,\n",
      "         0.3993271 , -0.3186609 , -0.09371175],\n",
      "       [-0.3134033 , -0.05702177, -0.0183657 , -0.55646414,  0.22632359,\n",
      "        -0.3674894 , -0.10589072, -0.6343746 ],\n",
      "       [-0.43294445, -0.41595346,  0.0602956 , -0.64939815, -0.12099046,\n",
      "        -0.43745494,  0.11576544,  0.31152278],\n",
      "       [-0.25505096, -0.39988703, -0.01722234,  0.13530765,  0.01034649,\n",
      "        -0.11702549, -0.22377518, -0.50821847],\n",
      "       [ 0.13459039, -0.5056072 ,  0.20428395, -0.23129769, -0.57473594,\n",
      "        -0.05845718,  0.48911893, -0.31750697]], dtype=float32), array([ 0.01810724,  0.01344353, -0.04203089,  0.07068869, -0.00339514,\n",
      "       -0.15687571,  0.00743723,  0.06446496], dtype=float32)), (), (array([[ 0.26193464],\n",
      "       [-0.1619365 ],\n",
      "       [-0.02677098],\n",
      "       [ 0.3586753 ],\n",
      "       [ 0.06419684],\n",
      "       [-0.88338447],\n",
      "       [ 0.20554483],\n",
      "       [ 0.8915742 ]], dtype=float32), array([0.497801], dtype=float32))]\n"
     ]
    }
   ],
   "source": [
    "params = sf.reveal(params_spu)\n",
    "print(params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, let's validate the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "auc=0.7982343165766513\n"
     ]
    }
   ],
   "source": [
    "auc = validate_model(params, X_test, y_test)\n",
    "print(f'auc={auc}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the end of the lab."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('sf')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "db45a4cb4cd37a8de684dfb7fcf899b68fccb8bd32d97c5ad13e5de1245c0986"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
